module dejan_accel
USE etienne_gino
use run_madx
use pointer_lattice   ! added
real(dp), allocatable :: rf_time(:),total_time(:),total_delta(:),my_tune(:,:)
type(fibre_array), allocatable ::cavity_dejan(:)
INTEGER :: NCAV_DEJAN=0, n_accel_dejan=0, INVERT_VOLT =1,mff
real(dp) :: n_atom , n_charge,n_ratio,DE
real(dp) :: delmin= -0.5d0, delmax=0.5d0
real(dp) :: kinmin,kinmax
real(dp) :: pmin,pmax
type(internal_state),pointer :: dejan_state
!logical :: invert_dejan=.true.
REAL(DP) AC(6,6),closed(6),OMEGA0
type(work) w0
real(dp), allocatable :: freqs(:),ts(:),x5(:),x6(:)
TYPE(fibre), POINTER :: pcav
TYPE(fibre), POINTER :: pt

TYPE PTCBEAM
   REAL(DP), POINTER :: X(:,:)
   LOGICAL(LP), POINTER :: U(:)
   INTEGER, POINTER :: N,LOST
END TYPE PTCBEAM
TYPE(PTCBEAM) B
TYPE(PTCBEAM) Bi

contains 



subroutine  set_distribution_for_beam(r1,nrays,nr,ndelta)  ! nrays is per ring, nr # of ring
IMPLICIT NONE
type(fibre), pointer :: p
TYPE(LAYOUT), target :: r1
integer i,n,J,nrays,nr,ndelta,k,mf,nt
REAL(DP) CIRC,X(6),ds,harm,domega,ddel,SIG(3),a(6,6),freq0,x0(6),thin,t
type(work) wi,wf,wt
LOGICAL(LP) U
TYPE(REAL_8) Y(6)
TYPE(NORMALFORM) NORMAL
TYPE(DAMAP) ID
type(internal_state) state
dejan_state=> default
call print(dejan_state,6)
!CAVITY_TOTALPATH=0

call kanalnummer(mf,"settings.dat")

!dejan_state=> my_estate

state=dejan_state+only_4d0

read(mf,*) NCAV_DEJAN,nt

allocate(ts(NCAV_DEJAN*nt),freqs(NCAV_DEJAN*nt))
allocate(x5(0:NCAV_DEJAN*nt),x6(0:NCAV_DEJAN*nt))
x5(0)=0.d0
x6(0)=0.d0
  DO K=1,NT*NCAV_DEJAN
     
	 read(mf,*)  n,J,freqs(k),ts(k),x5(k),x6(k)
	 
  ENDDO

close(mf)

x=0.d0
x(5)=my_delta 
call find_orbit(r1,x,my_start,state,1.d-5)
write(6,*) x
call init(state,1,0)
call alloc(id)
call alloc(y)
call alloc(NORMAL)
id=1
y=id+x

call track(r1,y,my_start,state)

NORMAL=y

write(6,*) NORMAL%tune
a=NORMAL%a_t
sig(1)=3.d-3
sig(2)=3.d-3
sig(3)=0.003d0
!sig=0.d0
write(6,*) " x max ", sqrt((a(1,1)**2+a(1,2)**2)*sig(1))
write(6,*) " y max ", sqrt((a(3,3)**2+a(3,4)**2)*sig(2))
call create_ptc_beam_target(B,nrays,nr,ndelta,SIG,A,x)


call kill(id)
call kill(y)
call kill(NORMAL)



end subroutine  set_distribution_for_beam

subroutine  plot_beam(final)
implicit none
integer i,mf
character(*),optional :: final
   call egm_scale_draw_ptcbeam(b,1)
   call egm_scale_draw_ptcbeam(b,2)
   call egm_scale_draw_ptcbeam(b,3)  !,invert_dejan)
   call kanalnummer(mf)
if(present(final)) then
 open(unit=mf,file=final)
else
 open(unit=mf,file='final.dat')
endif
DO I=1,b%n
  WRITE(mf,'(6(1x,E24.17),1x,l2)') B%X(I,1:6),b%u(i)
ENDDO
close(mf)

end subroutine plot_beam

subroutine  track_accel_dejan(L,t,u,x)
implicit none
TYPE(LAYOUT), POINTER :: L
type(integration_node), pointer :: t
logical(lp) u
real(dp) x(6)
integer n,i

pt=>pcav
do i=1,l%n
   ! write(6,*) pt%mag%name,pt%pos
    IF(pt%MAG%KIND==KIND4) THEN
     pt%MAG%C4%freq=freqs(n_accel_dejan)
     pt%MAG%C4%T=ts(n_accel_dejan)
     pt%MAGP%C4%freq=freqs(n_accel_dejan)
     pt%MAGP%C4%T=ts(n_accel_dejan)
	  if(.not.u) then
       CALL TRACK_probe_X(x,dejan_state, fibre1=pt, fibre2=pt%next)
u=.not.check_stable
t=>lost_node
	  endif
      pt=>pt%NEXT
	  exit
	ELSE
	  if(.not.u) then
        CALL TRACK_probe_X(x,dejan_state, fibre1=pt, fibre2=pt%next)
u=.not.check_stable
t=>lost_node
	  endif
      pt=>pt%NEXT
    ENDIF
	 
enddo
end subroutine  track_accel_dejan

subroutine  track_accel_dejan_beam(L)
implicit none
TYPE(LAYOUT), POINTER :: L
integer i
type(integration_node), pointer :: t
logical(lp) u
if(n_accel_dejan==0) pcav=>l%start
if(n_accel_dejan==size(ts)) return

     n_accel_dejan=n_accel_dejan+1
do i=1,b%n

if(.not.b%u(i)) then
 u=.false.
  call track_accel_dejan(L,t,u,b%x(i,:))
if(u) then
  CALL RESET_APERTURE_FLAG
  b%lost=b%lost+1
  write(6,*) B%lost, " lost in tracking routine "
  b%x(i,:)=bi%x(i,:)
 endif
 b%u(i)=u
endif

enddo

pcav=>pt

 if(mod(n_accel_dejan,ncav_dejan)==1) then
! if(n>0) then
   call egm_scale_draw_ptcbeam(b,1)
   call egm_scale_draw_ptcbeam(b,2)
   call egm_scale_draw_ptcbeam(b,3)  !,n,invert_dejan)
!  endif
 endif

end subroutine  track_accel_dejan_beam

  subroutine create_ptc_beam_target(B,nperring,nr,ndelta,SIG,A,C)
    USE gauss_dis
    implicit none
    INTEGER N,I,J,K,kk,nr,nperring,ja,kr,ndelta
    REAL(DP) SIG(3),X
    TYPE(PTCBEAM) B
    REAL(DP), OPTIONAL :: A(6,6),C(6)
    REAL(DP)  AS(6,6),CL(6),XT(6),r(3)
    
    n=nperring*(2*ndelta+1)*nr
    IF(.NOT.ASSOCIATED(B%N)) THEN
       CALL ALLOCATE_PTCBEAM(B,N)
    ELSEIF(B%N/=N) THEN
       CALL KILL_PTCBEAM(B)
       CALL ALLOCATE_PTCBEAM(B,N)
    ENDIF
    IF(.NOT.ASSOCIATED(Bi%N)) THEN
       CALL ALLOCATE_PTCBEAM(Bi,N)
    ELSEIF(Bi%N/=N) THEN
       CALL KILL_PTCBEAM(Bi)
       CALL ALLOCATE_PTCBEAM(Bi,N)
    ENDIF
    write(6,*) B%n, " RAYS "
    CL=ZERO; AS=ZERO;
    DO I=1,6
     AS(I,I)=ONE
    ENDDO

    IF(PRESENT(A)) AS=A
    IF(PRESENT(C)) CL=C
    r(1)=sqrt(sig(1))
    r(2)=sqrt(sig(2))
    r(3)=2*sig(3)/(2*ndelta+1)

    kr=0
    DO I=1,nr
		do ja=1,nperring
			DO J=1,2
				XT(2*J-1)= i*(r(j)*cos((twopi/nperring)*ja)/nr)
				XT(2*J)= i*(r(j)*sin((twopi/nperring)*ja)/nr)
			ENDDO
			XT(5)= 0.d0
			xt(6)=0.d0

			do kk=-ndelta,ndelta
				kr=kr+1
				B%X(kr,1:6)=CL(:)
				DO J=1,6
					DO K=1,6
						B%X(kr,J)=AS(J,K)*XT(K)+B%X(kr,J)
					ENDDO        
				ENDDO
				B%X(kr,5)=B%X(kr,5)+kk*r(3)
!				B%X(kr,6)=B%X(kr,6)+kk*r(3)
			enddo

		enddo
    ENDDO
  bi%x=b%x
   
  end    subroutine create_ptc_beam_target

  SUBROUTINE ALLOCATE_PTCBEAM(B,N)
    IMPLICIT NONE
    TYPE(PTCBEAM) , INTENT (INOUT) :: B
    INTEGER , INTENT (IN) :: N
    INTEGER I

    ALLOCATE(B%N,B%LOST)

    B%N=N
    B%LOST=0
    ALLOCATE(B%X(N,6))
    ALLOCATE(B%U(1:N))

    B%X  = ZERO
    B%U  = .FALSE.
  END SUBROUTINE ALLOCATE_PTCBEAM

  SUBROUTINE KILL_PTCBEAM(B)
    IMPLICIT NONE
    TYPE(PTCBEAM) , INTENT (INOUT) :: B
    IF(ASSOCIATED(B%N)) THEN
       DEALLOCATE(B%N,B%LOST,B%X,B%U)
    ENDIF
  END SUBROUTINE KILL_PTCBEAM

  subroutine stat_ptcbeam_raw(B,n,xm)
    implicit none
    INTEGER i,j,k,mf,NOTlost,N
    TYPE(PTCBEAM), INTENT(IN):: B
    TYPE(INTEGRATION_NODE),POINTER::T
    TYPE(FIBRE),POINTER::F
    real(dp), optional :: xm(6)
    real(dp), allocatable :: av(:,:),c(:)
    real(dp) em(3),beta(3),xma(6)
    mf=0
	if(mf==0) then
	 call kanalnummer(mf)
	 open(unit=mf,file='junkstat.txt')
	endif
    allocate(av(n,n),c(n))
    av=zero
    c=zero
    notlost=0
    xma(:)=-one
    DO K=1,b%n
       IF(.not.B%U(K)) THEN
          do i=1,6
            if(abs(b%x(k,i))>xma(i)) xma(i)=abs(b%x(k,i))
          enddo

          do i=1,n
             c(i)=b%x(k,i)+c(i)
             do j=i,n
                av(i,j)= b%x(k,i)*b%x(k,j)+av(i,j)
             enddo
          enddo
          notlost=notlost+1
       ENDIF
    ENDDO
    IF(NOTLOST==0) THEN
       if(mf/=6) then
          WRITE(mf,*) " ALL PARTICLES ARE LOST "
          WRITE(mf,*) " NO STATISTICS "
		  close(mf)
       else
          WRITE(6,*) " ALL PARTICLES ARE LOST "
          WRITE(6,*) " NO STATISTICS "
       endif
       deallocate(av)
       RETURN
    ENDIF
    if(notlost/=b%n-b%lost) then
       Write(6,*) " Error keeping track of lost particles "
       stop 999
    endif

    WRITE(MF,*) " NUMBER LEFT ",B%N-B%LOST
    if(mf/=6)WRITE(6,*) " NUMBER LEFT ",B%N-B%LOST
    WRITE(MF,*) " LOST ",B%LOST
    if(mf/=6)WRITE(6,*) " LOST ",B%LOST
    av=av/notlost
    c=c/notlost
    av(1,1)=av(1,1)-c(1)**2
    av(2,2)=av(2,2)-c(2)**2
    av(3,3)=av(3,3)-c(3)**2
    av(4,4)=av(4,4)-c(4)**2
    av(5,5)=av(5,5)-c(5)**2
    av(6,6)=av(6,6)-c(6)**2
    em(1)=two*sqrt(abs(av(1,1)*av(2,2)-av(1,2)**2))
    em(2)=two*sqrt(abs(av(3,3)*av(4,4)-av(3,4)**2))
    em(3)=two*sqrt(abs(av(5,5)*av(6,6)-av(5,6)**2))
    beta(1)=two*av(1,1)/em(1)
    beta(2)=two*av(3,3)/em(2)
    if(em(3)/=0.d0) beta(3)=two*av(5,5)/em(3)

    write(mf,*) " average arrays "
    write(mf,*) "betas ",beta
    write(mf,*) "emittances ",em
    if(mf/=6) then
       write(6,*) " average arrays "
       write(6,*) "betas ",beta
       write(6,*) "emittances ",em
    endif
    write(6,*) " limits "
    write(6,*) xma(1:2)
    write(6,*) xma(3:4)
    write(6,*) xma(5:6)
    if(present(xm)) xm=xma
100 FORMAT(7(1x,e13.6))
		  close(mf)

    deallocate(av)
  END subroutine stat_ptcbeam_raw


!   call egm_scale_draw_ptcbeam(b,1,.false.,deltak=total_delta(n))
!   call egm_scale_draw_ptcbeam(b,2,.false.,deltak=total_delta(n))
!   call egm_scale_draw_ptcbeam(b,3,invert_dejan,total_time(n),total_delta(n))

subroutine egm_scale_draw_PTCbeam(b,nd)  !,invert,phase,deltak)
implicit none 
integer nd,i,nt
real x,y 
type(PTCbeam) b
integer colour


if(nd==1) then
OSCIL_=.NOT.OSCIL_
  IF(OSCIL_) THEN
   colour=colour_stable 
  ELSE
   colour=gyellow
  ENDIF  

endif
 
call gmActivateGraphicsFrame( plot(nd))

call gmClearGraphicsFrame(plot(nd))

do i=1,b%n

if(b%u(i)) then
  x=bi%x(i,2*nd-1)  !-b%x(1,2*nd-1)
  y=bi%x(i,2*nd)    !-b%x(1,2*nd)
  if(nd==3.and.n_accel_dejan>0) then
	!if(i==1) write(6,*) "delta ",x
    x=x-x5(0)
    y=y-x6(0)
  endif
 call egm_scale_draw(plot(nd),x,y,gred)
else
  x=b%x(i,2*nd-1)  !-b%x(1,2*nd-1)
  y=b%x(i,2*nd)    !-b%x(1,2*nd)
  if(nd==3.and.n_accel_dejan>0) then
	!if(i==1) write(6,*) "delta ",x
    x=x-x5(n_accel_dejan)
    y=y-x6(n_accel_dejan)
  else
  if(x<2*plot_size(nd)%xmin) b%u(i)=.true.
  if(x>2*plot_size(nd)%xmax) b%u(i)=.true.
  if(y<2*plot_size(nd)%ymin) b%u(i)=.true.
  if(y>2*plot_size(nd)%ymax) b%u(i)=.true.
  endif
  if(b%u(i)) then
!   nt=nint(float(n)/ncav_dejan)
   b%lost=b%lost+1
   Write(6,*) " particle ",i,"lost in plane",nd
   Write(6,*) b%n-b%lost," particles still surviving  " 
   Write(6,*) x,y
   Write(6,*) plot_size(nd)
  ! write(6,*) my_tune(nt,1:3)
  ! Write(mff,*) " particle ",i,"lost in plane",nd
  !  Write(mff,*) x,y
  !write(mff,*) nt,my_tune(nt,1:3)
   !my_tune
  endif

! if(nd==3) then
!  write(6,*) invert,nd,b%x(i,2*nd-1),b%x(i,2*nd)
!  write(6,*) x,y,z
! pause 444
! endif

!if(abs(b%x(i,5)-dkk)<1.e-4) then
! call egm_scale_draw(plot(nd),x,y,gred)
!else
 call egm_scale_draw(plot(nd),x,y,colour)
!endif

endif
if(mod(i,100)==0) call gFlushGraphics

enddo
 call gFlushGraphics

end subroutine egm_scale_draw_PTCbeam

subroutine track_ideal(R)
  use run_madx
  implicit none
  type(layout), intent(inout), target :: R
  integer i,NT,K,LC,nc
  type(fibre), pointer :: p
  REAL(DP) CIRC,FREQ,VOLT,EG,X(6),EMAX,EMIN,Y(6)
  TYPE(WORK) W0,W
  
  DEFAULT=DEFAULT0+TIME0+TOTALPATH0 !+fringe0
  FEED_P0C=.FALSE.

  CALL GET_LENGTH(R,CIRC)
  VOLT=-5.D-3/48.D0
  EG=1.D-6/48.D0
  EMAX=817.D-6
  EMIN=218.D-6
  
  NT=NINT((EMAX-EMIN)/EG/48)
  WRITE(6,*) " NUMBER OF TURNS ",NT

  p=>R%START
  W0=P
  W=W0
  WRITE(6,*) W0%MASS
  WRITE(6,*) W0%BETA0
  WRITE(6,*) W0%P0C
  WRITE(6,*) W0%ENERGY
  WRITE(6,*) CIRC
!    O=twopi*EL%freq/CLIGHT
!       x(5)=x(5)-el%f(ko)*dir*EL%volt*c_1d_3*SIN(ko*O*(x(6)-EL%t)+EL%PHAS+ &
!            EL%PH(KO)+EL%phase0)/EL%P%P0C
  FREQ=4.D0*W0%BETA0*CLIGHT/CIRC
  WRITE(6,*) " FREQ ",FREQ
X=0.D0
Y=0.D0
nc=0
  p=>R%START
   do i=1,R%N
    IF(P%MAG%KIND==KIND4) THEN
	 P%MAG%C4%PHASE0=0.D0
	 P%MAG%C4%PHAS=0.D0
	 P%MAG%C4%VOLT=VOLT
	 P%MAGP%C4%PHASE0=0.D0
	 P%MAGP%C4%PHAS=0.D0
	 P%MAGP%C4%VOLT=VOLT
	 LC=I
	 nc=nc+1
	 ENDIF
     p=>P%NEXT
   ENDDO 
   
   Y(6)=0.1D0
OPEN(UNIT=20,FILE='CAV.DAT')
OPEN(UNIT=21,FILE='PLOT.DAT')
OPEN(UNIT=22,FILE='SETTINGS.DAT')
  write(22,*) nc,nt
  DO K=1,NT
  p=>R%START

  do i=1,R%N
    IF(P%MAG%KIND==KIND4) THEN
     P%MAG%C4%freq=4.D0*W%BETA0*CLIGHT/CIRC
     P%MAG%C4%T=X(6)+(1.D0/(twopi*P%MAG%C4%freq/CLIGHT))*ASIN(EG/VOLT/1.D-3)
     P%MAGP%C4%freq=P%MAG%C4%freq
     P%MAGP%C4%T=P%MAG%C4%T
     WRITE(22,'(I4,1X,I4,1X,4(E20.13,1X))') K,I,P%MAG%C4%freq,P%MAG%C4%T,x(5),x(6)
     CALL TRACK(R,X,I,I+1,DEFAULT)
     CALL TRACK(R,Y,I,I+1,DEFAULT)
     W=EG
     IF(I==LC) WRITE(20,'(I4,1X,2(E15.8,1X))') K,P%MAG%C4%freq,W0%ENERGY+X(5)*W0%P0C-W0%MASS
     IF(I==LC) WRITE(21,'(I4,1X,2(E15.8,1X))') K,Y(5),X(5)
	ELSE
     CALL TRACK(R,X,I,I+1,DEFAULT)
     CALL TRACK(R,Y,I,I+1,DEFAULT)
    ENDIF
	 
  p=>P%NEXT

  ENDDO
   WRITE(6,*) K,W0%ENERGY+X(5)*W0%P0C-W0%MASS
  ENDDO
 
 CLOSE(20)
 CLOSE(21)
 CLOSE(22)
end subroutine track_ideal 


subroutine build_emodel_rf(lattice)
  use run_madx
  implicit none
  type(layout), intent(out), target :: lattice

  type(layout) :: cell
  type(fibre) :: dl,dm,ds,fb,db,cv
  real(dp) :: kinetic,energy,brho,p0c
  real(dp) :: ang1,angf,angd
  real(dp) :: ldrl,ldrm,ldrs,lf,ld,gf,gd,kf,kd
  real(dp) :: lrf,frf,vrf,phs
  integer :: hrf

  call make_states(.true.)           ! use electrons
  exact_model = .true.                ! use full square-root Hamiltonian
  default = default + nocavity
  ALWAYS_EXACTMIS=.true.
  call update_states
  madlength = .false.                 ! use arc length, even for RBENDs

  !---------------------------------------------
  !  define the electron model lattice:
  !        ----====--========--====----|=====|
  !         dm  fb ds   db   ds fb  dm   cav
  ! L/cm ~   4  4.5 2    9    2 4.5  4    8
  !---------------------------------------------

  kinetic = 218.d-6                 ! kinetic energy of electrons /GeV
  call set_mad(kinetic = kinetic, method = 2, step = 20)
  call get_one(energy = energy)
  call get_one(brho = brho)
  call get_one(p0c = p0c)
  madkind2=drift_kick_drift

  ang1 = (twopi / 48.d0)                    ! angle of one cell
  angf = -1.0d0 * (twopi * 2.80d0 / 360.d0) ! angle of focusing bends
  angd = ang1 - 2.0d0 * angf                ! angle of defocusing bends
  lf = 0.0445d0                      ! arclength of focusing bends /m
  ld = 0.0890d0                      ! arclength of defocusing bends /m
  gf =  0.358d0                      ! focusing strength /T.m^-1
  gd = -0.314d0                      ! defocusing strength /T.m^-1
  kf = gf / brho                     ! focusing strength /m^-2
  kd = gd / brho                     ! defocusing strength /m^-2
  ldrl = 0.08055                     ! length of long drift /m
  ldrm = 0.04000                     ! length of medium drift /m
  ldrs = 0.01909                     ! length of short drift /m
  lrf  = 0.08110                     ! length of rf cavity /m

  hrf = 1        ! harmonic number
  frf = 0.d0     ! rf frequency (would like to ramp this, 47.7 --> 63.5 MHz, but not sure how to do this)
  vrf = 0.d-3    ! voltage across single rf cavity / MV (surely not 25d-6 ?!?)
  phs = 0.25d0   ! phase lag / 2pi (?)

  dl = drift("longdr",ldrl)
  dm = drift("medidr",ldrm)
  ds = drift("shrtdr",ldrs)
  fb = sbend("fbend",lf,angf).q.kf
  db = sbend("dbend",ld,angd).q.kd
!  cv = rfcavity("rfcav",lrf,vrf,phs,hrf,frf)
  cv = rfcavity("rfcav",0.d0,vrf,phs,hrf,frf)

  cell = dl + fb + ds + db + ds + fb + dl + cv
!  cell = dm + fb + ds + db + ds + fb + dm + cv
  lattice = 48 * cell
  lattice = .ring.lattice
  call survey(lattice)
end subroutine build_emodel_rf

end module dejan_accel
